package cn.edu.sjtu.stap.cfg.builder;

import java.util.*;

import org.aspectj.ajdt.internal.compiler.ast.*;
import org.aspectj.org.eclipse.jdt.internal.compiler.ast.*;
import org.aspectj.org.eclipse.jdt.internal.compiler.lookup.*;
import org.aspectj.weaver.AdviceKind;

import cn.edu.sjtu.stap.cfg.*;
import cn.edu.sjtu.stap.cg.*;
import cn.edu.sjtu.stap.graph.*;

/**
 * Inter-procedural CFG Builder. It builds inter-procedural CFG for the programs
 * according to the CG and intra-procedural CFGs made from the programs.
 * 
 * @author William
 * 
 */
public class InterModuleCfgBuilder {
    private Map<AbstractMethodDeclaration, ControlFlowGraph> interCfgs = null;

    private ControlFlowGraph currentCfg = null;

    public InterModuleCfgBuilder() {
        initialize();
    }

    public void initialize() {
        interCfgs = new HashMap<AbstractMethodDeclaration, ControlFlowGraph>();
        currentCfg = null;
    }

    public void buildInterModuleCfg(CallGraph cg, Map<AbstractMethodDeclaration, ControlFlowGraph> cfgs,
            AbstractMethodDeclaration main) {
        currentCfg = internalBuild(cg, cfgs, main);

        interCfgs.put(main, currentCfg);
    }

    private ControlFlowGraph internalBuild(CallGraph cg, Map<AbstractMethodDeclaration, ControlFlowGraph> cfgs,
            AbstractMethodDeclaration main) {
        ControlFlowGraph mainControl = cfgs.get(main);
        if (null == mainControl) {
            throw new IllegalArgumentException("main() not found!");
        }
        ControlFlowGraph result = new ControlFlowGraph(mainControl);

        Queue<AbstractMethodDeclaration> methodsToProcess = new LinkedList<AbstractMethodDeclaration>();
        Set<AbstractMethodDeclaration> processed = new HashSet<AbstractMethodDeclaration>();
        methodsToProcess.add(main);

        while (!methodsToProcess.isEmpty()) {
            // Get the method.
            AbstractMethodDeclaration callerMethod = methodsToProcess.poll();
            if (processed.contains(callerMethod)) {
                continue;
            }

            // Get the CFG of this method.
            ControlFlowGraph callerCfg = cfgs.get(callerMethod);

            // Traverse its nodes to connect the CFGs.
            Iterator<GraphNode> nIt = callerCfg.getNodes().iterator();
            while (nIt.hasNext()) {
                CfgNode node = (CfgNode) nIt.next();
                if (node.getKind() == CfgNodeKind.CALL && !(node.getASTNode() instanceof Proceed)) {
                    CfgCall callNode = (CfgCall) node;

                    connectCallSite(cg, cfgs, callerMethod, callNode, result, methodsToProcess);

                    weaveCallSite(cg, cfgs, callerMethod, callNode, result, methodsToProcess);
                }
            }

            weaveMethod(cg, cfgs, callerMethod, result, methodsToProcess);

            processed.add(callerMethod);
        }

        return result;
    }

    private void connectCallSite(CallGraph cg, Map<AbstractMethodDeclaration, ControlFlowGraph> cfgs,
            AbstractMethodDeclaration callerMethod, CfgCall callNode, ControlFlowGraph result,
            Queue<AbstractMethodDeclaration> methodsToProcess) {
        // Find all possible callees of this method.
        Set<AbstractMethodDeclaration> callees = getCallees(cg, callerMethod, callNode.getASTNode());
        Iterator<AbstractMethodDeclaration> calleeIt = callees.iterator();
        while (calleeIt.hasNext()) {
            AbstractMethodDeclaration method = calleeIt.next();
            ControlFlowGraph toAdd = cfgs.get(method);
            if (toAdd != null) {
                result.addControlFlowGraph(toAdd, callNode, callNode.getReturnNode(), CfgFlowKind.CALL_FLOW,
                        CfgFlowKind.RETURN_FLOW);
                methodsToProcess.add(method);
            }
        }
    }

    private void weaveCallSite(CallGraph cg, Map<AbstractMethodDeclaration, ControlFlowGraph> cfgs,
            AbstractMethodDeclaration callerMethod, CfgCall callNode, ControlFlowGraph result,
            Queue<AbstractMethodDeclaration> methodsToProcess) {
        ControlFlowGraph callerCfg = cfgs.get(callerMethod);

        // Find all advices that advises this call site.
        InvocationSite is = (InvocationSite) callNode.getASTNode();
        List<CgAdvice> advices = null;
        try {
            advices = cg.findAdvices(is);

            if (!advices.isEmpty()) {
                // System.out.println("Advices in order on:\n" + is + "\n: " +
                // advices + "\n");
            }

            // If the above statement does not throw any exceptions,
            // it goes here.

            Map<CfgCall, ControlFlowGraph> callSitesMap = new HashMap<CfgCall, ControlFlowGraph>();
            // For multiple proceed() in around advice.
            Map<CfgCall, ControlFlowGraph> newlyAddedCallSitesMap = new HashMap<CfgCall, ControlFlowGraph>();

            callSitesMap.put(callNode, callerCfg);
            Iterator<CgAdvice> i = advices.iterator();
            while (i.hasNext()) {
                CgAdvice a = i.next();
                AdviceDeclaration advice = (AdviceDeclaration) a.getASTNode();
                AdviceKind k = advice.kind;

                ControlFlowGraph adviceCfg = cfgs.get(advice);

                boolean goneIntoAround = false;
                boolean adviceProcessed = false;
                // Apply the advice on all call sites.
                Iterator<CfgCall> callIt = callSitesMap.keySet().iterator();
                while (callIt.hasNext()) {
                    CfgCall currentCall = callIt.next();
                    CfgReturn currentReturn = currentCall.getReturnNode();
                    ControlFlowGraph cfgToWeave = callSitesMap.get(currentCall);

                    if (k == AdviceKind.Before) {
                        weaveAdviceBefore(currentCall, adviceCfg, cfgToWeave, result);
                    } else if (k == AdviceKind.Around) {

                        // Weave the advice's CFG into current
                        // method's CFG.
                        weaveAdviceExternalAround(currentCall, currentReturn, cfgToWeave, adviceCfg, result);

                        // Add current method's CFG into result.
                        // result.add(cfgToWeave);
                        // result.remove(currentCall);
                        // result.remove(currentReturn);

                        if (!adviceProcessed) {
                            Iterator<GraphNode> nodeIt = adviceCfg.getNodes().iterator();
                            while (nodeIt.hasNext()) {
                                CfgNode n = (CfgNode) nodeIt.next();
                                if (n.getKind() == CfgNodeKind.CALL && n.getASTNode() instanceof Proceed) {
                                    CfgCall currentCallNode = (CfgCall) n;
                                    CfgReturn currentReturnNode = currentCallNode.getReturnNode();

                                    // Reconnect all the possible callees with
                                    // proceed().
                                    Set<AbstractMethodDeclaration> callees = getCallees(cg, callerMethod, callNode
                                            .getASTNode());
                                    Iterator<AbstractMethodDeclaration> calleeIt = callees.iterator();
                                    while (calleeIt.hasNext()) {
                                        AbstractMethodDeclaration method = calleeIt.next();
                                        ControlFlowGraph toAdd = cfgs.get(method);
                                        if (toAdd != null) {
                                            result.addControlFlowGraph(toAdd, currentCallNode, currentReturnNode,
                                                    CfgFlowKind.CALL_FLOW, CfgFlowKind.RETURN_FLOW);
                                        }
                                    }
                                    newlyAddedCallSitesMap.put(currentCallNode, adviceCfg);
                                }
                            }
                            adviceProcessed = true;
                        }
                        goneIntoAround = true;

                    } else if (k == AdviceKind.After) {
                        weaveAdviceAfter(currentReturn, adviceCfg, cfgToWeave, result);
                    } else if (k == AdviceKind.AfterReturning) {
                        weaveAdviceAfter(currentReturn, adviceCfg, cfgToWeave, result);
                    } else if (k == AdviceKind.AfterThrowing) {
                        // Not concerned yet.
                    }
                }

                if (goneIntoAround) {
                    callSitesMap.clear();
                    callSitesMap.putAll(newlyAddedCallSitesMap);
                    newlyAddedCallSitesMap.clear();
                }

                methodsToProcess.add(advice);
            }
        } catch (CallSiteNotFoundException e) {
            // Skip quitely.
        }
    }

    private void weaveMethod(CallGraph cg, Map<AbstractMethodDeclaration, ControlFlowGraph> cfgs,
            AbstractMethodDeclaration callerMethod, ControlFlowGraph result,
            Queue<AbstractMethodDeclaration> methodsToProcess) {
        ControlFlowGraph callerCfg = cfgs.get(callerMethod);

        // Weave for methods.
        List<CgAdvice> advices = null;
        try {
            advices = cg.findAdvices(callerMethod);

            if (!advices.isEmpty()) {
                // System.out.println("Advices in order on:\n" + callerMethod +
                // "\n: " + advices + "\n");
            }

            // If the above statement does not throw any exceptions,
            // it goes here.
            Iterator<CgAdvice> i = advices.iterator();
            ControlFlowGraph cfgToWeave = callerCfg;
            CfgNode beforeInsertion = cfgToWeave.getCfgEntry();
            CfgNode afterInsertion = cfgToWeave.getCfgReturningExit();
            while (i.hasNext()) {
                CgAdvice a = i.next();
                AdviceDeclaration advice = (AdviceDeclaration) a.getASTNode();
                AdviceKind k = advice.kind;
                ControlFlowGraph adviceCfg = cfgs.get(advice);

                if (k == AdviceKind.Before) {
                    beforeInsertion = weaveAdviceAfter(beforeInsertion, adviceCfg, cfgToWeave, result);
                } else if (k == AdviceKind.Around) {
                    ControlFlowGraph newCfg = weaveAdviceInternalAround(beforeInsertion, afterInsertion, callerCfg,
                            adviceCfg, result);
                    // result.add(callerMethodGraph);
                    AbstractMethodDeclaration artificialMethod = new ArtificialMethodDeclaration(callerMethod);
                    cfgs.put(artificialMethod, newCfg);

                    // Connect proceed() with this new cfg.
                    Iterator<GraphNode> nodeIt = adviceCfg.getNodes().iterator();
                    while (nodeIt.hasNext()) {
                        CfgNode n = (CfgNode) nodeIt.next();
                        if (n.getKind() == CfgNodeKind.CALL && n.getASTNode() instanceof Proceed) {
                            CfgCall callNode = (CfgCall) n;
                            CfgReturn returnNode = callNode.getReturnNode();
                            // In fact here, should add to
                            // callMethodGraph
                            // first, then result.
                            result.addControlFlowGraph(newCfg, callNode, returnNode, CfgFlowKind.CALL_FLOW,
                                    CfgFlowKind.RETURN_FLOW);
                        }
                    }

                    methodsToProcess.add(artificialMethod);

                    // Reset before insertion and after insertion.
                    beforeInsertion = newCfg.getCfgEntry();
                    afterInsertion = newCfg.getCfgReturningExit();
                } else if (k == AdviceKind.After) {
                    afterInsertion = weaveAdviceBefore(afterInsertion, adviceCfg, cfgToWeave, result);
                } else if (k == AdviceKind.AfterReturning) {
                    afterInsertion = weaveAdviceBefore(afterInsertion, adviceCfg, cfgToWeave, result);
                } else if (k == AdviceKind.AfterThrowing) {
                    // Not concerned yet.
                }

                methodsToProcess.add(advice);
            }
        } catch (MethodNotFoundException e) {
            // Skip quietly.
        }
    }

    private CfgNode weaveAdviceBefore(CfgNode node, ControlFlowGraph advice, ControlFlowGraph callerCfg,
            ControlFlowGraph result) {
        CfgWeave weaveNode = new CfgWeave();
        CfgWeaveReturn weaveReturnNode = new CfgWeaveReturn();

        reconnectIncomingEdges(node, weaveNode);

        CfgFlow f = new CfgFlow(CfgFlowKind.NORMAL_FLOW, weaveReturnNode, node);
        callerCfg.add(f);

        result.add(f);
        result.addControlFlowGraph(advice, weaveNode, weaveReturnNode, CfgFlowKind.WEAVE_FLOW,
                CfgFlowKind.WEAVE_RETURN_FLOW);

        return weaveNode;
    }

    private void weaveAdviceExternalAround(CfgNode callNode, CfgNode returnNode, ControlFlowGraph callerCfg,
            ControlFlowGraph adviceCfg, ControlFlowGraph result) {
        CfgWeave weaveNode = new CfgWeave();
        CfgWeaveReturn weaveReturnNode = new CfgWeaveReturn();

        reconnectIncomingEdges(callNode, weaveNode);
        reconnectOutgoingEdges(returnNode, weaveReturnNode);

        callerCfg.remove(callNode);
        callerCfg.remove(returnNode);

        result.remove(callNode);
        result.remove(returnNode);
        result.addControlFlowGraph(adviceCfg, weaveNode, weaveReturnNode, CfgFlowKind.WEAVE_FLOW,
                CfgFlowKind.WEAVE_RETURN_FLOW);
    }

    private ControlFlowGraph weaveAdviceInternalAround(CfgNode entry, CfgNode exit, ControlFlowGraph callerCfg,
            ControlFlowGraph adviceCfg, ControlFlowGraph result) {
        CfgWeave weaveNode = new CfgWeave();
        CfgWeaveReturn weaveReturnNode = new CfgWeaveReturn();

        // Keep the original content of the method for further use.
        ControlFlowGraph newCfg = new ControlFlowGraph(ControlFlowGraph.ARTIFICIAL_METHOD_CFG);
        newCfg.removeFlow(newCfg.getCfgEntry(), newCfg.getCfgReturningExit());

        // Remove the nodes and edges between entry and exit node, and move them
        // into the newly created method.
        moveNodes(callerCfg, result, entry, exit, newCfg);

        CfgFlow f1 = new CfgFlow(CfgFlowKind.NORMAL_FLOW, entry, weaveNode);
        CfgFlow f2 = new CfgFlow(CfgFlowKind.NORMAL_FLOW, weaveReturnNode, exit);
        callerCfg.add(f1);
        callerCfg.add(f2);

        result.add(f1);
        result.add(f2);
        result.addControlFlowGraph(adviceCfg, weaveNode, weaveReturnNode, CfgFlowKind.WEAVE_FLOW,
                CfgFlowKind.WEAVE_RETURN_FLOW);

        return newCfg;
    }

    private CfgNode weaveAdviceAfter(CfgNode node, ControlFlowGraph advice, ControlFlowGraph callerCfg,
            ControlFlowGraph result) {
        CfgWeave weaveNode = new CfgWeave();
        CfgWeaveReturn weaveReturnNode = new CfgWeaveReturn();

        reconnectOutgoingEdges(node, weaveReturnNode);

        CfgFlow f = new CfgFlow(CfgFlowKind.NORMAL_FLOW, node, weaveNode);

        callerCfg.add(f);

        result.add(f);
        result.addControlFlowGraph(advice, weaveNode, weaveReturnNode, CfgFlowKind.WEAVE_FLOW,
                CfgFlowKind.WEAVE_RETURN_FLOW);

        return weaveReturnNode;
    }

    public void reconnectIncomingEdges(CfgNode original, CfgNode current) {
        if (null == original) {
            throw new NullPointerException("Original node null!");
        }
        if (null == current) {
            throw new NullPointerException("Current node null!");
        }

        Iterator<GraphEdge> i = original.getIncomingEdgeSet().iterator();
        List<GraphEdge> toChange = new ArrayList<GraphEdge>();
        while (i.hasNext()) {
            toChange.add(i.next());
        }
        i = toChange.iterator();
        while (i.hasNext()) {
            GraphEdge edge = i.next();
            edge.setDestinationNode(current);
        }
    }

    public void reconnectOutgoingEdges(CfgNode original, CfgNode current) {
        if (null == original) {
            throw new NullPointerException("Original node null!");
        }
        if (null == current) {
            throw new NullPointerException("Current node null!");
        }

        Iterator<GraphEdge> i = original.getOutgoingEdgeSet().iterator();
        List<GraphEdge> toChange = new ArrayList<GraphEdge>();
        while (i.hasNext()) {
            toChange.add(i.next());
        }
        i = toChange.iterator();
        while (i.hasNext()) {
            GraphEdge edge = i.next();
            edge.setSourceNode(current);
        }
    }

    private boolean addIndependentNode(Queue<CfgNode> independentNodeList, Set<CfgNode> processedNodeList, CfgNode node) {
        if (processedNodeList.contains(node)) {
            return false;
        }
        // Iterator<GraphEdge> i = node.getIncomingEdgeSet().iterator();
        // while (i.hasNext()) {
        // CfgFlow f = (CfgFlow) i.next();
        // if (f.getKind() != CfgFlowKind.LOOP_FLOW) {
        // return false;
        // }
        // }

        independentNodeList.add(node);
        return true;
    }

    private void moveNodes(ControlFlowGraph src, ControlFlowGraph result, CfgNode startNode, CfgNode endNode,
            ControlFlowGraph dest) {
        Queue<CfgNode> independentNodes = new LinkedList<CfgNode>();
        Set<CfgNode> processedNodes = new HashSet<CfgNode>();

        // Remove and transfer the outgoing edges and destination nodes of
        // startNode.
        GraphEdgeSet outgoingEdges = new GraphEdgeSet();
        outgoingEdges.addAll(startNode.getOutgoingEdgeSet());
        Iterator<GraphEdge> eIt = outgoingEdges.iterator();
        while (eIt.hasNext()) {
            CfgFlow edge = (CfgFlow) eIt.next();
            CfgNode node = edge.getDestinationNode();

            // Remove from the original places.
            src.remove(edge);
            result.remove(edge);

            addIndependentNode(independentNodes, processedNodes, node);

            // Transfer to the new place.
            edge.setSourceNode(dest.getCfgEntry());
            if (node == endNode) {
                edge.setDestinationNode(dest.getCfgReturningExit());
            } else {
                edge.setDestinationNode(node);
            }
            dest.add(edge);
        }

        // Remove and transfer the rest edges and nodes.
        while (!independentNodes.isEmpty()) {
            CfgNode sourceNode = independentNodes.poll();

            if (sourceNode == endNode) {
                continue;
            }

            if (sourceNode.getKind() == CfgNodeKind.CALL) {
                dest.add(sourceNode);

                independentNodes.add(((CfgCall) sourceNode).getReturnNode());
            } else {
                outgoingEdges.clear();
                outgoingEdges.addAll(sourceNode.getOutgoingEdgeSet());
                Iterator<GraphEdge> i = outgoingEdges.iterator();
                while (i.hasNext()) {
                    CfgFlow edge = (CfgFlow) i.next();
                    CfgNode destinationNode = edge.getDestinationNode();

                    // Remove from the original places.
                    src.remove(edge);
                    result.remove(edge);

                    addIndependentNode(independentNodes, processedNodes, destinationNode);

                    // Transfer to the new place.
                    edge.setSourceNode(sourceNode);
                    if (destinationNode == endNode) {
                        edge.setDestinationNode(dest.getCfgReturningExit());
                    } else {
                        edge.setDestinationNode(destinationNode);
                    }
                    dest.add(edge);
                }
            }
            processedNodes.add(sourceNode);
        }

        Iterator<CfgNode> nIt = processedNodes.iterator();
        while (nIt.hasNext()) {
            CfgNode toRemoveNode = nIt.next();

            src.remove(toRemoveNode);
            result.remove(toRemoveNode);
        }
    }

    private Set<AbstractMethodDeclaration> getCallees(CallGraph cg, AbstractMethodDeclaration caller, ASTNode methodCall) {
        MethodBinding b = null;

        if (methodCall instanceof MessageSend) {
            MessageSend messageSend = (MessageSend) methodCall;
            b = messageSend.binding;
        } else if (methodCall instanceof ExplicitConstructorCall) {
            ExplicitConstructorCall explicitConstructor = (ExplicitConstructorCall) methodCall;
            b = explicitConstructor.binding;
        } else if (methodCall instanceof AllocationExpression) {
            AllocationExpression allocationExpression = (AllocationExpression) methodCall;
            b = allocationExpression.binding;
        }

        Set<AbstractMethodDeclaration> result = new HashSet<AbstractMethodDeclaration>();
        if (b != null) {
            AbstractMethodDeclaration callee = b.sourceMethod();
            List<CgMethod> callees = null;
            try {
                if (caller instanceof ArtificialMethodDeclaration) {
                    caller = ArtificialMethodDeclaration.getMethod((ArtificialMethodDeclaration) caller);
                }
                callees = cg.findOverridingCallees(caller, callee);
            } catch (MethodNotFoundException e) {
                System.err.println("Error occurred while finding callees!");
                e.printStackTrace();
            }
            if (callees != null) {
                Iterator<CgMethod> i = callees.iterator();
                while (i.hasNext()) {
                    CgMethod node = i.next();
                    if (node.getASTNode() != null) {
                        result.add((AbstractMethodDeclaration) node.getASTNode());
                    }
                }
            }
        }

        return result;
    }

    public Map<AbstractMethodDeclaration, ControlFlowGraph> getResult() {
        return interCfgs;
    }
}

class ArtificialMethodDeclaration extends MethodDeclaration {
    private static HashMap<AbstractMethodDeclaration, ArtificialMethodDeclaration> map = new HashMap<AbstractMethodDeclaration, ArtificialMethodDeclaration>();

    public ArtificialMethodDeclaration(AbstractMethodDeclaration method) {
        super(null);
        map.put(method, this);
    }

    public static AbstractMethodDeclaration getMethod(ArtificialMethodDeclaration m) {
        Iterator<AbstractMethodDeclaration> i = map.keySet().iterator();
        while (i.hasNext()) {
            AbstractMethodDeclaration method = i.next();
            if (map.get(method) == m) {
                return method;
            }
        }
        return null;
    }
}